#!/usr/bin/env python
'''
This module provides functions to read in the input
from PerM's output ( *.mapping file ) and generate the
output for TranSEQ, which estimate the abundance level. 
, given the target gene name.
Created on Aug 22, 2009
@author: yanghoch
'''
from PerM_output_format import *
from common import print_list
# A simple matrix
# This matrix is a list of lists
# Column and row numbers start with 1

class Matrix( object ):
    def __init__( self, rows, cols ):
        self.cols = cols
        self.rows = rows
        # initialize matrix and fill with zeroes
        self.matrix = []
        for i in range( rows ):
            ea_row = []
            for j in range( cols ):
                ea_row.append( 0 )
            self.matrix.append( ea_row )

    def setitem( self, row, col, v ):
        self.matrix[row][col] = v

    def getitem( self, row, col ):
        return self.matrix[row][col]

    def __repr__( self ):
        outStr = ""
        for i in range( self.rows ):
            outStr += 'Row %s = %s\n' % ( i + 1, self.matrix[i] )
        return outStr

    def is_empty_col( self, colId ):
        colIsEmpty = True
        for i in range( 0, self.rows ):
            if( self.getitem( i, colId ) > 0 ):
                colIsEmpty = False
                break
        return colIsEmpty

class MappingsOfIsoforms2AGene( object ):
    def __init__( self, geneName ):
        DEFAULT_READ_LENGTH = 50
        self.geneName = geneName
        self.isoforms = []
        self.reads = []
        self.mappings = [] # list of list. Each read has a list for
        self.read_length = DEFAULT_READ_LENGTH

    def add_map( self, read, isoform ):
        if read in self.reads:
            readId = self.reads.index( read )
        else:
            readId = len( self.reads )
            self.reads.append( read )
        if isoform in self.isoforms:
            isoformId = self.isoforms.index( isoform )
        else:
            isoformId = len( self.isoforms )
            self.isoforms.append( isoform )
        if ( len( self.mappings ) <= readId ):
            self.mappings.append( [] )
        self.mappings[readId].append( isoformId )

    def add_mapping( self, geneName, read, isoform ):
        if( geneName == self.geneName ):
            self.add_map( read, isoform )
        else:
            pass

    def add_mappings_from_a_file( self, mapping_file ):
        ''' read a *.mapping file of a gene from PerM 
        and add the mapping info to self.mappings DS'''
        # TODO complete the function 
        ifile = open( mapping_file, mode = "r" )
        for line in ifile:
            mapping = PerMMappingLineInfo( line )
            self.add_map( mapping.read_name, mapping.isoform_name )
            del mapping

    def __fill_hit_matrix__( self, isoform_lens ):
        '''fill the weighted-hit matrix'''
        rown = len( self.reads )
        coln = len( self.isoforms )
        m = Matrix( rown, coln )
        for r in range( 0, rown ):
            for i in range( 0, coln ):
                if i in self.mappings[r]: # this isoform is mapped by the read r
                    # lookup dictionary
                    isoform_len = isoform_lens[self.isoforms[i]]
                    denominator = max( isoform_len - self.read_length, 1 )
                    # Assume the mapped read number
                    # is proportional to the denominator  
                    value = m.getitem( r, i ) + 1.0 / float( denominator )
                    m.setitem( r, i, value )
        return( m )

    def get_TranSEQ_input( self, isoform_length_dict ):
        ''' Given a dict of isoform vs it's length, 
        this function output a matrix with 
        isoform num (column num) x reads num (rows num). 
        The value of each cell is the hit num / isoform length'''

        if( len( isoform_length_dict ) < len( self.isoforms ) ):
            print len( isoform_length_dict )
            print " < " + str( len( self.isoforms ) )
            print "The number of isoform are inconsistent!"
            print_list ( self.isoforms )
            return( -1 )
        # used the list of list 
        # and divide the cell values with isoforms' lengths 
        matrix = self.__fill_hit_matrix__( isoform_length_dict )
        ofile = open( self.geneName + ".dat", mode = "w" )
        # output the isoforms name as the header
        for i in range( 0, matrix.cols ):
            ofile.write( self.isoforms[i] )
            if( i < matrix.cols - 1 ):
                ofile.write( '\t' );
        ofile.write( '\n' )

        for i in range( 0, matrix.rows ):
            ofile.write( self.reads[i] )
            for j in range( 0, matrix.cols ):
                ofile.write( '\t' )
                if j in self.mappings[i]:
                    ofile.write( str( matrix.getitem( i, j ) ) )
                else:
                    ofile.write( str( 0 ) )
            ofile.write( str( '\n' ) )
        ofile.close()

